package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"
)

// 定义配置结构
type Config struct {
	IP     string `json:"ip"`
	Port   string `json:"port"`
	APIKey string `json:"api_key"`
}

// 定义数据结构
type DatasetResponse struct {
	Data    []Dataset `json:"data"`
	HasMore bool      `json:"has_more"`
	Limit   int       `json:"limit"`
	Total   int       `json:"total"`
	Page    int       `json:"page"`
}

type Dataset struct {
	ID            string `json:"id"`
	Name          string `json:"name"`
	Description   string `json:"description"`
	UpdatedAt     int64  `json:"updated_at"`
	Tags          []Tag  `json:"tags"`
	DocumentCount int    `json:"document_count"`
	WordCount     int    `json:"word_count"`
}

type Tag struct {
	ID   string `json:"id"`
	Name string `json:"name"`
	Type string `json:"type"`
}

type DocumentResponse struct {
	Data    []Document `json:"data"`
	HasMore bool       `json:"has_more"`
	Limit   int        `json:"limit"`
	Total   int        `json:"total"`
	Page    int        `json:"page"`
}

type Document struct {
	ID          string `json:"id"`
	Position    int    `json:"position"`
	DatasetID   string `json:"dataset_id"`
	Name        string `json:"name"`
	CreatedAt   int64  `json:"created_at"`
	CompletedAt int64  `json:"completed_at"`
	IndexingAt  int64  `json:"indexing_at"`
	Enabled     bool   `json:"enabled"`
	DisabledAt  int64  `json:"disabled_at"`
	DisabledBy  string `json:"disabled_by"`
	HitCount    int    `json:"hit_count"`
	WordCount   int    `json:"word_count"`
}

type SegmentResponse struct {
	Data    []Segment `json:"data"`
	HasMore bool      `json:"has_more"`
	Limit   int       `json:"limit"`
	Total   int       `json:"total"`
	Page    int       `json:"page"`
}

type Segment struct {
	ID          string `json:"id"`
	Position    int    `json:"position"`
	DocumentID  string `json:"document_id"`
	Content     string `json:"content"`
	WordCount   int    `json:"word_count"`
	Tokens      int    `json:"tokens"`
	IndexingAt  int64  `json:"indexing_at"`
	CompletedAt int64  `json:"completed_at"`
	Enabled     bool   `json:"enabled"`
	DisabledAt  int64  `json:"disabled_at"`
	DisabledBy  string `json:"disabled_by"`
	HitCount    int    `json:"hit_count"`
	CreatedBy   string `json:"created_by"`
	CreatedAt   int64  `json:"created_at"`
}

const configFileName = "dify_config.json"

func main() {
	fmt.Println("=== Dify知识库下载工具 ===")

	// 1. 尝试读取配置文件，如果不存在则获取用户输入
	config, err := loadConfig()
	if err != nil {
		if os.IsNotExist(err) {
			fmt.Println("首次使用，需要配置Dify平台信息")
			config = getConfigFromUser()
			if err := saveConfig(config); err != nil {
				fmt.Printf("警告: 保存配置失败: %v\n", err)
			} else {
				fmt.Println("配置已保存，下次启动将自动使用")
			}
		} else {
			fmt.Printf("读取配置失败: %v\n", err)
			return
		}
	} else {
		fmt.Printf("已加载配置: %s:%s，如需修改，请编辑配置文件\n", config.IP, config.Port)
	}

	baseURL := fmt.Sprintf("http://%s:%s", config.IP, config.Port)

	// 2. 获取知识库列表
	datasets, err := getDatasets(baseURL, config.APIKey)
	if err != nil {
		fmt.Printf("获取知识库列表失败: %v\n", err)
		// 如果API调用失败，可能是配置已过期，询问是否重新配置
		if askForReconfigure() {
			config = getConfigFromUser()
			if err := saveConfig(config); err != nil {
				fmt.Printf("警告: 保存配置失败: %v\n", err)
			}
			// 重新尝试获取知识库列表
			baseURL = fmt.Sprintf("http://%s:%s", config.IP, config.Port)
			datasets, err = getDatasets(baseURL, config.APIKey)
			if err != nil {
				fmt.Printf("再次获取知识库列表失败: %v\n", err)
				return
			}
		} else {
			return
		}
	}

	// 3. 显示知识库列表供用户选择
	selectedDataset := selectDataset(datasets)
	if selectedDataset == nil {
		fmt.Println("未选择知识库，程序退出")
		return
	}

	// 4. 获取分段分隔符设置，并自动保存
	separator := getSegmentSeparator()

	// 5. 获取选定知识库的文档列表
	documents, err := getDocuments(baseURL, config.APIKey, selectedDataset.ID)
	if err != nil {
		fmt.Printf("获取文档列表失败: %v\n", err)
		return
	}

	// 6. 下载知识库内容
	err = downloadKnowledgeBase(baseURL, config.APIKey, selectedDataset, documents, separator)
	if err != nil {
		fmt.Printf("下载知识库失败: %v\n", err)
		return
	}

	fmt.Println("知识库下载完成！")
}

func loadConfig() (*Config, error) {
	// 获取可执行文件所在目录
	exePath, err := os.Executable()
	if err != nil {
		return nil, err
	}
	exeDir := filepath.Dir(exePath)
	configPath := filepath.Join(exeDir, configFileName)

	fmt.Println("加载配置文件：" + configPath)
	data, err := os.ReadFile(configPath)
	if err != nil {
		return nil, err
	}

	var config Config
	err = json.Unmarshal(data, &config)
	if err != nil {
		return nil, err
	}

	return &config, nil
}

func saveConfig(config *Config) error {
	// 获取可执行文件所在目录
	exePath, err := os.Executable()
	if err != nil {
		return err
	}
	exeDir := filepath.Dir(exePath)
	configPath := filepath.Join(exeDir, configFileName)

	data, err := json.MarshalIndent(config, "", "  ")
	if err != nil {
		return err
	}

	return os.WriteFile(configPath, data, 0644)
}

func getConfigFromUser() *Config {
	reader := bufio.NewReader(os.Stdin)
	config := &Config{}

	fmt.Print("请输入Dify平台IP地址: ")
	ip, _ := reader.ReadString('\n')
	config.IP = strings.TrimSpace(ip)

	if config.IP == "" {
		fmt.Println("IP地址不能为空！")
		os.Exit(1)
	}

	fmt.Print("请输入端口号(默认80): ")
	port, _ := reader.ReadString('\n')
	config.Port = strings.TrimSpace(port)
	if config.Port == "" {
		config.Port = "80"
	}

	fmt.Print("请输入API Key: ")
	apiKey, _ := reader.ReadString('\n')
	config.APIKey = strings.TrimSpace(apiKey)

	if config.APIKey == "" {
		fmt.Println("API Key不能为空！")
		os.Exit(1)
	}

	return config
}

func askForReconfigure() bool {
	reader := bufio.NewReader(os.Stdin)
	fmt.Print("API调用失败，是否重新配置连接信息？(y/n, 默认n): ")
	answer, _ := reader.ReadString('\n')
	answer = strings.TrimSpace(strings.ToLower(answer))
	return answer == "y" || answer == "yes"
}

func getSegmentSeparator() string {
	reader := bufio.NewReader(os.Stdin)

	fmt.Println("\n=== 分段分隔符设置 ===")
	fmt.Println("请选择分段之间的分隔符（导入时需要设置相同的分隔符）：")
	fmt.Println("1. ---- (4个减号)")
	fmt.Println("2. ==== (4个等号，默认)")
	fmt.Println("3. **** (4个星号)")
	fmt.Println("4. 自定义分隔符")
	fmt.Print("请选择(1-4，默认2): ")

	choice, _ := reader.ReadString('\n')
	choice = strings.TrimSpace(choice)

	switch choice {
	case "1":
		return "----"
	case "2", "":
		return "===="
	case "3":
		return "****"
	case "4":
		fmt.Print("请输入自定义分隔符: ")
		customSep, _ := reader.ReadString('\n')
		customSep = strings.TrimSpace(customSep)
		if customSep == "" {
			fmt.Println("使用默认分隔符 ====")
			return "===="
		}
		return customSep
	default:
		fmt.Println("无效选择，使用默认分隔符 ====")
		return "===="
	}
}

func getDatasets(baseURL, apiKey string) ([]Dataset, error) {
	url := fmt.Sprintf("%s/v1/datasets?page=1&limit=50", baseURL)

	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}

	req.Header.Set("Authorization", "Bearer "+apiKey)
	req.Header.Set("Content-Type", "application/json")

	client := &http.Client{Timeout: 30 * time.Second}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(body))
	}

	var datasetResponse DatasetResponse
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	err = json.Unmarshal(body, &datasetResponse)
	if err != nil {
		return nil, err
	}

	return datasetResponse.Data, nil
}

func selectDataset(datasets []Dataset) *Dataset {
	fmt.Println("\n=== 知识库列表 ===")
	fmt.Printf("%-3s %-40s %-20s %-10s %-8s %s\n",
		"序号", "名称", "更新时间", "文档数", "字数", "标签")
	fmt.Println(strings.Repeat("-", 100))

	for i, dataset := range datasets {
		updateTime := time.Unix(dataset.UpdatedAt, 0).Format("2006-01-02 15:04:05")

		// 处理标签
		tagNames := make([]string, len(dataset.Tags))
		for j, tag := range dataset.Tags {
			tagNames[j] = tag.Name
		}
		tagsStr := strings.Join(tagNames, ",")
		if tagsStr == "" {
			tagsStr = "无"
		}

		fmt.Printf("%-3d %-40s %-20s %-10d %-8d %s\n",
			i+1,
			truncateString(dataset.Name, 38),
			updateTime,
			dataset.DocumentCount,
			dataset.WordCount,
			truncateString(tagsStr, 20))
	}

	reader := bufio.NewReader(os.Stdin)
	fmt.Print("\n请选择知识库(输入序号，输入0退出): ")
	choiceStr, _ := reader.ReadString('\n')
	choiceStr = strings.TrimSpace(choiceStr)

	choice, err := strconv.Atoi(choiceStr)
	if err != nil || choice < 0 || choice > len(datasets) {
		fmt.Println("无效的选择！")
		return nil
	}

	if choice == 0 {
		return nil
	}

	return &datasets[choice-1]
}

func getDocuments(baseURL, apiKey, datasetID string) ([]Document, error) {
	url := fmt.Sprintf("%s/v1/datasets/%s/documents?page=1&limit=100", baseURL, datasetID)

	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}

	req.Header.Set("Authorization", "Bearer "+apiKey)
	req.Header.Set("Content-Type", "application/json")

	client := &http.Client{Timeout: 30 * time.Second}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(body))
	}

	var docResponse DocumentResponse
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	err = json.Unmarshal(body, &docResponse)
	if err != nil {
		return nil, err
	}

	return docResponse.Data, nil
}

func getSegments(baseURL, apiKey, datasetID, documentID string) ([]Segment, error) {
	url := fmt.Sprintf("%s/v1/datasets/%s/documents/%s/segments?page=1&limit=1000", baseURL, datasetID, documentID)

	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}

	req.Header.Set("Authorization", "Bearer "+apiKey)
	req.Header.Set("Content-Type", "application/json")

	client := &http.Client{Timeout: 30 * time.Second}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(body))
	}

	var segResponse SegmentResponse
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	err = json.Unmarshal(body, &segResponse)
	if err != nil {
		return nil, err
	}

	return segResponse.Data, nil
}

func downloadKnowledgeBase(baseURL, apiKey string, dataset *Dataset, documents []Document, separator string) error {
	// 创建目录
	dirName := fmt.Sprintf("knowledge_base_%s_%s",
		strings.ReplaceAll(dataset.Name, "/", "_"),
		time.Now().Format("20060102_150405"))

	err := os.MkdirAll(dirName, 0755)
	if err != nil {
		return fmt.Errorf("创建目录失败: %v", err)
	}

	fmt.Printf("正在下载知识库到目录: %s\n", dirName)
	fmt.Printf("分段分隔符: %s\n", separator)

	// 保存知识库基本信息和分隔符设置
	datasetInfo := map[string]interface{}{
		"name":              dataset.Name,
		"description":       dataset.Description,
		"document_count":    dataset.DocumentCount,
		"word_count":        dataset.WordCount,
		"updated_at":        dataset.UpdatedAt,
		"download_time":     time.Now().Format("2006-01-02 15:04:05"),
		"segment_separator": separator,
		"import_note":       fmt.Sprintf("导入时请设置分段分隔符为: %s", separator),
	}

	datasetInfoJSON, _ := json.MarshalIndent(datasetInfo, "", "  ")
	err = os.WriteFile(dirName+"/dataset_info.json", datasetInfoJSON, 0644)
	if err != nil {
		return fmt.Errorf("保存知识库信息失败: %v", err)
	}

	// 下载每个文档的内容
	totalSegments := 0
	for i, doc := range documents {
		fmt.Printf("正在下载文档 %d/%d: %s\n", i+1, len(documents), doc.Name)

		segments, err := getSegments(baseURL, apiKey, dataset.ID, doc.ID)
		if err != nil {
			fmt.Printf("警告: 获取文档 %s 的分段失败: %v\n", doc.Name, err)
			continue
		}

		// 保存文档内容
		filename := fmt.Sprintf("%s/%02d_%s.txt", dirName, i+1,
			strings.ReplaceAll(doc.Name, "/", "_"))

		file, err := os.Create(filename)
		if err != nil {
			fmt.Printf("警告: 创建文件 %s 失败: %v\n", filename, err)
			continue
		}

		// 写入每个分段内容
		for _, segment := range segments {
			// 使用用户选择的分隔符
			file.WriteString(separator + "\n")
			file.WriteString(segment.Content)
			file.WriteString("\n\n")
		}

		file.Close()
		totalSegments += len(segments)

		fmt.Printf("  文档 %s 下载完成，包含 %d 个分段\n", doc.Name, len(segments))
	}

	// 创建导入说明文件
	createImportGuide(dirName, separator)

	fmt.Printf("\n下载完成！总计下载 %d 个文档，%d 个分段\n", len(documents), totalSegments)
	fmt.Printf("文件保存在: %s/\n", dirName)
	fmt.Printf("分段分隔符: %s (请记住此分隔符，导入时需要设置)\n", separator)

	return nil
}

func createImportGuide(dirName, separator string) {
	guide := fmt.Sprintf(`# 知识库导入指南

## 重要提示
在导入此知识库时，请确保设置正确的分段分隔符。

## 分段分隔符设置
分段分隔符: %s

## 导入步骤
1. 在Dify平台创建新的知识库
2. 上传文档时，在"分段设置"中选择"自定义分隔符"
3. 在分隔符输入框中填写: %s
4. 上传下载的文档文件

## 文件说明
- dataset_info.json: 知识库基本信息
- *.txt: 各个文档的内容文件
- IMPORT_GUIDE.md: 本导入指南

## 注意事项
- 确保使用正确的分隔符，否则分段可能不正确
- 建议先上传一个文档测试分段效果
- 如有问题，请检查分隔符设置是否正确`, separator, separator)

	err := os.WriteFile(dirName+"/IMPORT_GUIDE.md", []byte(guide), 0644)
	if err != nil {
		fmt.Printf("警告: 创建导入指南失败: %v\n", err)
	}
}

func truncateString(s string, length int) string {
	if len(s) <= length {
		return s
	}
	return s[:length-3] + "..."
}
